mod my_crypto;
mod param;

use std::{
    io::{BufWriter, Write},
    thread, vec,
};

use chrono::prelude::*;
use crossbeam_channel::{bounded, unbounded};
use mysql::prelude::*;
use oracle::{sql_type::OracleType, SqlValue};

#[derive(Debug)]
enum Handler {
    MySQL(mysql::Conn),
    File(BufWriter<std::fs::File>),
}

impl Handler {
    fn run_sql(&mut self, sql: &str, flag: u8, reset: bool, nolog: bool) -> bool {
        match self {
            Handler::File(f) => {
                f.write_all(sql.as_bytes()).unwrap();
                if flag == 0 {
                    // just run sql
                    f.write_all(";\n".as_bytes()).unwrap();
                } else if flag == 1 {
                    // commit and start new transaction
                    f.write_all(";\ncommit;\nbegin;\n".as_bytes()).unwrap();
                } else {
                    // commit
                    f.write_all(";\nSET FOREIGN_KEY_CHECKS = 1;\n".as_bytes())
                        .unwrap();

                    if nolog {
                        f.write_all("SET SQL_LOG_BIN = 1;\n".as_bytes()).unwrap();
                    }
                }
                f.flush().unwrap();
            }
            Handler::MySQL(m) => {
                match m.exec_drop(sql, ()) {
                    Ok(_) => {}
                    Err(e) => {
                        print_log(&format!("Insert failed: {}", e.to_string()));
                        print_log(&sql[..1024]);
                        return false;
                    }
                }
                if flag == 1 {
                    // commit and start new transaction
                    m.exec_drop("commit", ()).unwrap();
                    if reset {
                        m.reset().unwrap();
                    }

                    // m.exec_drop("SET AUTOCOMMIT = 0", ()).unwrap();
                    // m.exec_drop("SET FOREIGN_KEY_CHECKS = 0", ()).unwrap();
                }
            }
        }
        true
    }
}

const SQL_SIZE: usize = 32 * 1024 * 1024;
const CHANNEL_SIZE: usize = 20;

fn main() {
    let params = match param::Opt::init() {
        Some(p) => p,
        None => return,
    };

    let mut handlers = vec![];
    for _ in 0..params.thread {
        let handler = if params.online {
            let opc = mysql::OptsBuilder::new()
                .pass(params.tk.as_ref())
                .user(Some(&params.tu))
                .db_name(params.tb.as_ref())
                .ip_or_hostname(params.th.as_ref())
                .tcp_port(params.tp);
            let mut m = mysql::Conn::new(opc).unwrap();

            if params.trans > 1 {
                m.start_transaction(mysql::TxOpts::default()).unwrap();
            }
            if params.no_log {
                m.exec_drop("SET SQL_LOG_BIN = 0\0", ()).unwrap();
            }
            m.exec_drop("SET FOREIGN_KEY_CHECKS = 0", ()).unwrap();
            if params.trans > 1 {
                m.exec_drop("set AUTOCOMMIT = 0", ()).unwrap();
            }
            Handler::MySQL(m)
        } else {
            let file = std::fs::File::create(&params.file).unwrap();
            let mut buf = BufWriter::new(file);

            if params.no_log {
                buf.write_all("SET SQL_LOG_BIN = 0;\n".as_bytes()).unwrap();
            }
            buf.write_all("SET FOREIGN_KEY_CHECKS = 0;\n".as_bytes())
                .unwrap();
            if params.trans > 1 {
                buf.write_all("begin;\n".as_bytes()).unwrap();
            }
            buf.flush().unwrap();
            Handler::File(buf)
        };
        handlers.push(handler);
    }

    let (tx, rx) = bounded::<String>(CHANNEL_SIZE);
    let (ctx, crx) = unbounded::<bool>();

    let sub_thread = thread::spawn(move || {
        let oc = match oracle::Connection::connect(
            &params.ou,
            &params.ok,
            &format!("{}:{}/{}", params.oh, params.op, params.ob),
        ) {
            Ok(c) => c,
            Err(e) => {
                print_log(&format!(
                    "Connect to oracle database failed, {}",
                    e.to_string()
                ));
                return;
            }
        };

        if let Some(sql) = params.pre_sql {
            oc.execute(&sql, &[]).unwrap();
        }

        print_log("Querying data");
        let res = match oc.query(&params.sql, &[]) {
            Ok(n) => n,
            Err(e) => {
                print_log(&format!("sql: {} \n error: {}", &params.sql, e.to_string()));
                return;
            }
        };
        print_log("Query data completed");

        let mut cols = String::new();
        let mut col_type = vec![];
        for x in res.column_info() {
            cols.push_str(&format!("`{}`", x.name()));
            cols.push(',');
            match x.oracle_type() {
                // 0: number, 1: bianry, 2: clob, 20: string
                OracleType::Number(_, _) => col_type.push(0),
                OracleType::Int64 => col_type.push(0),
                OracleType::Float(_) => col_type.push(0),
                OracleType::BLOB => col_type.push(1),
                OracleType::Raw(_) => col_type.push(1),
                OracleType::BFILE => col_type.push(1),
                OracleType::LongRaw => col_type.push(1),
                OracleType::CLOB => col_type.push(2),
                _ => col_type.push(20),
            }
        }
        cols.pop();

        let mut sql_cmd = String::with_capacity(SQL_SIZE);

        sql_cmd.push_str(&format!("insert into `{}`({}) values (", params.tt, cols));

        let mut i = 0; // number rows in sql
        let mut count = 0; // total rows

        for val in res {
            for (j, x) in val.unwrap().sql_values().iter().enumerate() {
                let t = get_value(x, col_type[j]);
                if col_type[j] == 20 {
                    sql_cmd.push('\'');
                }
                sql_cmd.push_str(&t);
                if col_type[j] == 20 {
                    sql_cmd.push('\'');
                }
                sql_cmd.push(',');
            }
            sql_cmd.pop();
            sql_cmd.push_str("),(");
            i += 1;
            if i >= params.row {
                sql_cmd.pop();
                sql_cmd.pop();
                tx.send(sql_cmd).unwrap();
                sql_cmd = String::with_capacity(SQL_SIZE);

                count += i;
                i = 0;
                print_log(&format!("Complete {} rows", count));
                sql_cmd.push_str(&format!("insert into `{}`({}) values (", params.tt, cols));
            }

            if let Ok(true) = crx.try_recv() {
                print_log("Got error, sub thread quit");
                return;
            }
        }

        if i > 0 {
            sql_cmd.pop();
            sql_cmd.pop();
            tx.send(sql_cmd).unwrap();
            count += i;
        }

        for _ in 0..params.thread {
            tx.send(String::from("commit")).unwrap();
        }

        print_log(&format!("Completed, total {} rows", count));
    });

    let mut thread_pool = scoped_threadpool::Pool::new(params.thread);
    thread_pool.scoped(|scope| {
        for mut handler in handlers {
            let rx = rx.clone();
            scope.execute(move || {
                let mut t = 0; // number sql in transaction
                loop {
                    let re = rx.recv().unwrap();
                    if &re == "commit" {
                        handler.run_sql(&re, 2, params.reset, params.no_log);
                        return;
                    }

                    t += 1;
                    if params.trans > 1 && t >= params.trans {
                        if !handler.run_sql(&re, 1, params.reset, params.no_log) {
                            return;
                        }
                        t = 0;
                    } else {
                        if !handler.run_sql(&re, 0, params.reset, params.no_log) {
                            return;
                        }
                    }
                }
            });
        }
    });

    ctx.send(true).unwrap();
    sub_thread.join().unwrap();
}

fn get_value(v: &SqlValue, t: usize) -> String {
    match t {
        // 0: number, 1: bianry, 2: clob, 20: string
        0 => {
            if let Ok(temp) = v.get::<String>() {
                return temp;
            }
        }
        1 => {
            if let Ok(temp) = v.get::<Vec<u8>>() {
                return to_hex(&temp);
            }
        }
        2 => {
            if let Ok(temp) = v.get::<String>() {
                return to_hex(temp.as_bytes());
            }
        }
        _ => {
            if let Ok(temp) = v.get::<String>() {
                let mut buff = String::with_capacity(temp.len() * 2);
                for x in temp.chars() {
                    match x {
                        '\n' => buff.push_str("\\n"),
                        '\r' => buff.push_str("\\r"),
                        '"' => buff.push_str("\\\""),
                        '\'' => buff.push_str("\\'"),
                        '\\' => buff.push_str("\\\\"),
                        _ => buff.push(x),
                    }
                }

                return buff;
            }
        }
    }
    String::from("NULL")
}

static CHARS: &'static [u8] = b"0123456789ABCDEF";
fn to_hex(data: &[u8]) -> String {
    let mut v = Vec::with_capacity(data.len() * 2 + 2);
    v.push(48); // '0'
    v.push(120); // 'x'
    for d in data {
        v.push(CHARS[(d >> 4) as usize]);
        v.push(CHARS[(d & 0xf) as usize]);
    }
    unsafe { String::from_utf8_unchecked(v) } // sure is a valid utf8 string
}

fn print_log(log: &str) {
    let now = Local::now();
    println!("[{}]: {}", now.format("%Y-%m-%d %H:%M:%S"), log);
}
